import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import tushare as ts
from datetime import datetime, timedelta

# 初始化Tushare pro接口
pro = ts.pro_api('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')


class TianShanAluminumStrategy:
    def __init__(self):
        self.data = None
        self.positions = 0  # 持仓数量
        self.capital = 1000000  # 初始资金100万
        self.trade_log = []
        self.current_date = None
        self.target_stock = '002532.SZ'  # 天山铝业

    def fetch_data(self, start_date, end_date, freq='D'):
        """获取股票数据"""
        print(f"正在获取天山铝业({self.target_stock}) {freq}线数据...")

        if freq == 'D':  # 日线数据
            df = pro.daily(**{
                "ts_code": self.target_stock,
                "start_date": start_date,
                "end_date": end_date
            }, fields=[
                "ts_code", "trade_date", "open", "high", "low", "close",
                "pre_close", "change", "pct_chg", "vol", "amount"
            ])
        elif freq == 'W':  # 周线数据
            df = pro.weekly(**{
                "ts_code": self.target_stock,
                "start_date": start_date,
                "end_date": end_date
            }, fields=[
                "ts_code", "trade_date", "open", "high", "low", "close",
                "pre_close", "change", "pct_chg", "vol", "amount"
            ])
        elif freq == '60min':  # 60分钟线
            df = pro.pro_bar(**{
                "ts_code": self.target_stock,
                "start_date": start_date,
                "end_date": end_date,
                "freq": '60min'
            }, fields=[
                "ts_code", "trade_time", "open", "high", "low", "close",
                "pre_close", "change", "pct_chg", "vol", "amount"
            ])
            df.rename(columns={'trade_time': 'trade_date'}, inplace=True)

        df['trade_date'] = pd.to_datetime(df['trade_date'])
        df = df.sort_values('trade_date')

        # 计算技术指标
        df['ma5'] = df['close'].rolling(5).mean()
        df['ma20'] = df['close'].rolling(20).mean()
        df['rsi'] = self.calculate_rsi(df, 14)
        df['atr'] = self.calculate_atr(df, 14)  # 平均真实波幅，用于止损

        self.data = df
        print(f"数据获取完成，共获取{len(df)}条{freq}线记录")
        return df

    def calculate_rsi(self, df, window=14):
        """计算RSI指标"""
        delta = df['close'].diff()
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)

        avg_gain = gain.rolling(window).mean()
        avg_loss = loss.rolling(window).mean()

        rs = avg_gain / avg_loss
        return 100 - (100 / (1 + rs))

    def calculate_atr(self, df, window=14):
        """计算ATR(平均真实波幅)"""
        high_low = df['high'] - df['low']
        high_close = (df['high'] - df['close'].shift()).abs()
        low_close = (df['low'] - df['close'].shift()).abs()

        tr = pd.concat([high_low, high_close, low_close], axis=1).max(axis=1)
        return tr.rolling(window).mean()

    def generate_signals(self):
        """生成交易信号"""
        print("正在生成交易信号...")
        df = self.data.copy()

        # 双均线策略 + RSI + 突破策略
        df['signal'] = 0

        # 买入信号条件
        buy_condition = (
                (df['ma5'] > df['ma20']) &  # 5日均线上穿20日均线
                (df['rsi'] > 30) & (df['rsi'] < 70) &  # RSI在合理区间
                (df['close'] > df['high'].rolling(5).max().shift(1))  # 突破5日高点
        )

        # 卖出信号条件
        sell_condition = (
                (df['ma5'] < df['ma20']) |  # 5日均线下穿20日均线
                (df['rsi'] > 70) |  # RSI超买
                (df['close'] < df['low'].rolling(5).min().shift(1))  # 跌破5日低点
        )

        df.loc[buy_condition, 'signal'] = 1
        df.loc[sell_condition, 'signal'] = -1

        self.data = df
        return df

    def execute_trades(self, date):
        """执行交易"""
        self.current_date = date
        daily_data = self.data[self.data['trade_date'] == date]

        if len(daily_data) == 0:
            return

        signal = daily_data['signal'].iloc[0]
        price = daily_data['close'].iloc[0]
        atr = daily_data['atr'].iloc[0]

        # 卖出信号处理
        if signal == -1 and self.positions > 0:
            self.sell_stock(price)

        # 买入信号处理
        elif signal == 1 and self.positions == 0:
            # 使用ATR确定头寸规模 (风险控制)
            risk_per_trade = 0.02  # 每笔交易风险2%
            shares = (self.capital * risk_per_trade) / (1.5 * atr)
            shares = int(shares // 100) * 100  # 按100股整数倍
            if shares > 0:
                self.buy_stock(price, shares)

    def buy_stock(self, price, shares):
        """买入股票"""
        cost = price * shares
        self.capital -= cost
        self.positions = shares

        self.trade_log.append({
            'date': self.current_date,
            'action': 'buy',
            'price': price,
            'shares': shares,
            'capital': self.capital,
            'portfolio_value': self.get_portfolio_value(price),
            'stop_loss': price - 1.5 * self.data[self.data['trade_date'] == self.current_date]['atr'].iloc[0]
        })
        print(f"{self.current_date.date()} 买入 {self.target_stock} {shares}股 @ {price:.2f}")

    def sell_stock(self, price):
        """卖出股票"""
        proceeds = price * self.positions
        self.capital += proceeds

        entry_price = self.trade_log[-1]['price'] if self.trade_log else price
        profit = (price - entry_price) * self.positions
        profit_pct = (price / entry_price - 1) * 100

        self.trade_log.append({
            'date': self.current_date,
            'action': 'sell',
            'price': price,
            'shares': self.positions,
            'profit': profit,
            'profit_pct': profit_pct,
            'capital': self.capital,
            'portfolio_value': self.get_portfolio_value(price)
        })
        print(f"{self.current_date.date()} 卖出 {self.target_stock} {self.positions}股 @ {price:.2f} "
              f"盈利: {profit:.2f} ({profit_pct:.2f}%)")

        self.positions = 0

    def check_stop_loss_and_take_profit(self):
        """检查止损和止盈"""
        if self.positions == 0:
            return

        current_data = self.data[self.data['trade_date'] == self.current_date]
        if len(current_data) == 0:
            return

        price = current_data['close'].iloc[0]
        last_trade = [t for t in self.trade_log if t['action'] == 'buy'][-1]
        stop_loss = last_trade.get('stop_loss', price * 0.95)

        # 止损逻辑
        if price <= stop_loss:
            self.sell_stock(price)
        # 止盈逻辑 (2:1的风险回报比)
        elif price >= last_trade['price'] + 2 * (last_trade['price'] - stop_loss):
            self.sell_stock(price)

    def get_portfolio_value(self, current_price=None):
        """计算投资组合总价值"""
        if current_price is None:
            current_data = self.data[self.data['trade_date'] == self.current_date]
            current_price = current_data['close'].iloc[0] if len(current_data) > 0 else 0
        return self.capital + (current_price * self.positions if self.positions > 0 else 0)

    def backtest(self, start_date, end_date, freq='D'):
        """回测策略"""
        self.fetch_data(start_date, end_date, freq)
        self.generate_signals()

        date_range = pd.date_range(start=start_date, end=end_date)
        portfolio_values = []

        for date in date_range:
            if date not in self.data['trade_date'].values:
                continue

            self.current_date = date
            self.execute_trades(date)
            self.check_stop_loss_and_take_profit()

            # 记录每日投资组合价值
            portfolio_value = self.get_portfolio_value()
            portfolio_values.append({
                'date': date,
                'value': portfolio_value,
                'price': self.data[self.data['trade_date'] == date]['close'].iloc[0]
            })

        # 计算回测结果
        trade_log_df = pd.DataFrame(self.trade_log)
        portfolio_df = pd.DataFrame(portfolio_values)

        # 计算基准收益 (买入持有)
        initial_price = portfolio_df['price'].iloc[0]
        final_price = portfolio_df['price'].iloc[-1]
        benchmark_return = (final_price / initial_price - 1) * 100

        # 策略收益
        strategy_return = (portfolio_df['value'].iloc[-1] / 1000000 - 1) * 100

        print("\n回测结果:")
        print(f"初始资金: 1,000,000")
        print(f"最终价值: {portfolio_df['value'].iloc[-1]:,.2f}")
        print(f"策略收益率: {strategy_return:.2f}%")
        print(f"买入持有收益率: {benchmark_return:.2f}%")
        print(f"总交易次数: {len(trade_log_df) // 2}次完整买卖")

        if len(trade_log_df) > 0:
            winning_trades = trade_log_df[trade_log_df['action'] == 'sell']
            if len(winning_trades) > 0:
                win_rate = len(winning_trades[winning_trades['profit'] > 0]) / len(winning_trades) * 100
                avg_profit = winning_trades['profit_pct'].mean()
                print(f"胜率: {win_rate:.2f}%")
                print(f"平均每笔交易收益: {avg_profit:.2f}%")

        # 可视化
        self.visualize_results(portfolio_df, benchmark_return)

        return trade_log_df, portfolio_df

    def visualize_results(self, portfolio_df, benchmark_return):
        """可视化回测结果"""
        plt.figure(figsize=(16, 8))

        # 创建双轴图表
        ax1 = plt.gca()
        ax2 = ax1.twinx()

        # 投资组合价值曲线
        ax1.plot(portfolio_df['date'], portfolio_df['value'],
                 label='策略净值', color='blue', linewidth=2)

        # 基准曲线 (买入持有)
        ax1.plot(portfolio_df['date'],
                 portfolio_df['price'] * (1000000 / portfolio_df['price'].iloc[0]),
                 '--', label='买入持有净值', color='green', alpha=0.7)

        # 股价曲线
        ax2.plot(portfolio_df['date'], portfolio_df['price'],
                 label='股价', color='red', alpha=0.5)

        # 标记买卖点
        trades = pd.DataFrame(self.trade_log)
        if len(trades) > 0:
            buy_dates = trades[trades['action'] == 'buy']['date']
            buy_prices = [portfolio_df[portfolio_df['date'] == d]['price'].values[0] for d in buy_dates]
            ax2.scatter(buy_dates, buy_prices, color='green', marker='^', s=100, label='买入')

            sell_dates = trades[trades['action'] == 'sell']['date']
            sell_prices = [portfolio_df[portfolio_df['date'] == d]['price'].values[0] for d in sell_dates]
            ax2.scatter(sell_dates, sell_prices, color='red', marker='v', s=100, label='卖出')

        # 设置图表属性
        ax1.set_title('天山铝业(002532.SZ)交易策略回测结果', fontsize=15)
        ax1.set_xlabel('日期', fontsize=12)
        ax1.set_ylabel('净值', fontsize=12)
        ax2.set_ylabel('股价', fontsize=12)

        # 合并图例
        lines1, labels1 = ax1.get_legend_handles_labels()
        lines2, labels2 = ax2.get_legend_handles_labels()
        ax1.legend(lines1 + lines2, labels1 + labels2, loc='upper left')

        ax1.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.show()


# 使用示例
if __name__ == "__main__":
    strategy = TianShanAluminumStrategy()

    # 设置回测日期范围
    end_date = datetime.now().strftime('%Y%m%d')
    start_date = (datetime.now() - timedelta(days=365)).strftime('%Y%m%d')  # 1年数据

    # 运行回测 (可以更改freq为'D'(日线), 'W'(周线)或'60min'(60分钟线))
    trade_log, portfolio_values = strategy.backtest(start_date, end_date, freq='D')